package com.carp.shard.plugin.strategy;

import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.carp.shard.plugin.utils.IdGenerator;
import lombok.NonNull;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.reflection.MetaObject;
import org.springframework.stereotype.Component;

import java.util.Objects;

/**
 * 分片策略（db_ysb_order.ts_wholesale_buy_car）
 *
 * @author lile
 * @since 2024/5/10
 */
@Component
public class BuyCarShardStrategy extends AbstractShardStrategy {

    /**
     * 表名
     */
    public static final String TABLE_NAME = "db_ysb_order.ts_wholesale_buy_car";
    /**
     * 分片数量
     */
    public static final int SHARD_NUMS = 8;
    /**
     * 主键
     */
    private static final String PRIMARY_KEY = "id";
    /**
     * 分片键（取MyBatis参数名）
     */
    public static final String SHARD_KEY = "userId";

    @Override
    public boolean replaceTableName(@NonNull SQLExprTableSource x, @NonNull BoundSql boundSql) {
        x.setExpr(getShardTableName(boundSql.getParameterObject()));
        return true;
    }

    @Override
    public void processParams(@NonNull MetaObject metaObject, @NonNull BoundSql boundSql, @NonNull SchemaStatVisitor statVisitor) {
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
        if (!Objects.equals(SqlCommandType.INSERT, mappedStatement.getSqlCommandType())) {
            return;
        }
        setColumnValue(boundSql, PRIMARY_KEY, statVisitor.getColumns(), IdGenerator::getSnowflakeId);
    }

    @Override
    protected String getOriginalTableName() {
        return TABLE_NAME;
    }

    @Override
    public String getShardTableName(@NonNull Object parameterObject) {
        int userId = (int) getValue(parameterObject, SHARD_KEY);
        return TABLE_NAME + "_" + getShardByModFor2Power(userId);
    }

    /**
     * 模运算获取分片号（当分片数量为2的幂次方时）
     * 模运算优化为位运算提高性能
     *
     * @param key 分片键
     * @return 分片号
     */
    private int getShardByModFor2Power(int key) {
        return key & (SHARD_NUMS - 1);
    }

}
